import argparse
import sys
import os
import time
from deepsignal3.utils.process_utils import get_motif_seqs
from deepsignal3.utils.process_utils import complement_seq
from deepsignal3.utils.ref_reader import get_contig2len,get_contig2len_n_seq

neg_num = 0  # Replacing mp.Value since we're now single-threaded

def read_and_process_feature(input_file, output_file, ref_path, batch_size=1):
    print('Begin processing', flush=True)
    fea_list = []
    #motif_seqs = get_motif_seqs('CG', True)
    chrom2len,contigs = get_contig2len_n_seq(ref_path)
    #contigs_complement = complement_seq(contigs)

    with open(input_file, 'r') as f, open(output_file, 'w') as wf:
        for line in f:
            words = line.strip().split('\t')
            bisulfite = [float(x) for x in words[15].split(",")]
            seq=words[3]
            pred_pos = [int(x) for x in words[5].split(",")]
            ref_pos = [int(x) for x in words[16].split(",")]
            rpos=[]
            bs=[]
            chr=words[1]
            refseq=[]
            if len(ref_pos)<=1:
                continue
            filtered_list = [x for x in ref_pos if x != -1]
            if len(filtered_list)<=1:
                continue
            # if ref_pos[1]!=-1 and ref_pos[0]!=-1 and ref_pos[1]>ref_pos[0]:
            #     strand='+'
            # elif ref_pos[1]!=-1 and ref_pos[0]!=-1 and ref_pos[1]<ref_pos[0]:
            #     strand='-'
            # else:
            #     continue
            if  filtered_list[1]>filtered_list[0]:
                strand='+'
            elif  filtered_list[1]<filtered_list[0]:
                strand='-'
            else:
                continue
            flag=[]
            for i in range(len(pred_pos)):
                ppos=pred_pos[i]
                if seq[ppos-20:ppos+21]=='ATCCCAGCACTTTGGGAGGCCGAGGCGGGCGGATCACGAGG':
                    if ref_pos[i]==-1:
                        continue
                    rpos.append(ref_pos[i])
                    bs.append(bisulfite[i])
                    if strand=='-':
                        refseq.append(complement_seq(contigs[chr][ref_pos[i]-20:ref_pos[i]+21]))
                    else:
                        refseq.append(contigs[chr][ref_pos[i]-20:ref_pos[i]+21])
                    if refseq[-1]=='ATCCCAGCACTTTGGGAGGCCGAGGCGGGCGGATCACGAGG':
                        flag.append(1)
                    else:
                        flag.append(0)

            if len(rpos)==0:
                continue
            
            # ref_pos_text = ','.join(
            #             [str(x) for x in rpos])
            
            

            #bisulfite_text = ','.join(bs)
            for i in range(len(refseq)):
                fea_str = '\t'.join([chr,str(rpos[i]),str(bs[i]),refseq[i],str(flag[i]),strand])


            
                #fea_str = '\t'.join(words[:5]+[pred_pos_text] +words[6:15]+ [bisulfite_text] + [ref_pos_text])
                fea_list.append(fea_str)

                if len(fea_list) >= batch_size:
                    for one_features_str in fea_list:
                        wf.write(one_features_str + "\n")
                    wf.flush()
                    fea_list = []

        if len(fea_list) > 0:
            for one_features_str in fea_list:
                wf.write(one_features_str + "\n")
            wf.flush()

    print('Processing completed. Added negative number is: {}'.format(neg_num), flush=True)

def parse_args():
    parser = argparse.ArgumentParser("")
    parser.add_argument("--input", type=str, required=True)
    parser.add_argument("--output", type=str, required=True)
    parser.add_argument("--ref",default='/homeb/hpc/refs/genome/chm13v2.0.fa',type=str,required=False)
    return parser.parse_args()

def main():
    args = parse_args()
    
    read_and_process_feature(args.input, args.output, args.ref)

if __name__ == '__main__':
    sys.exit(main())
